import pathlib
import time

import cv2
import imutils
import numpy as np
from PIL import Image
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from script.ShopCenter import putCenter

model_id = 'damo/cv_vitb16_segmentation_shop-seg'
shop_seg = pipeline(Tasks.shop_segmentation, model=model_id)
base_path = pathlib.Path(__file__).resolve().parent.parent
image_matting_origin = pathlib.Path.joinpath(base_path, "image_matting_origin")
image_matting_done = pathlib.Path.joinpath(base_path, "image_matting_done")
pathlib.Path.mkdir(image_matting_origin, exist_ok=True)
pathlib.Path.mkdir(image_matting_done, exist_ok=True)


class Mask2Image:
    """抠图转换成png 含Alpha通道"""

    def __init__(self, mask: np.ndarray, image: np.ndarray):
        self.mask = mask
        self.image = image
        self.preprocess()

    @classmethod
    def read_filePath(cls, filePath):
        return cv2.imdecode(np.fromfile(filePath, dtype=np.uint8))

    def preprocess(self):
        """图片进行预处理, 腐蚀5个像素， 避免分割线突兀"""
        kernel = np.ones((5, 5), np.uint8)
        self.mask = cv2.erode(self.mask, kernel, iterations=1)

    def transform_png(self) -> np.ndarray:
        self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2BGRA)
        return cv2.bitwise_and(self.image, self.image, mask=self.mask)


class ReplaceBackground:
    """替换图片的背景"""

    def __init__(self, background_color=None, size=None, scale=None):
        self.background_color = background_color or (255, 255, 255)  # 背景默认白色
        self.size = size or (800, 800)  # 默认800X800
        self.scale = scale or 1  # 占比默认100%

    @classmethod
    def cv2PIL(cls, image):
        """opencv和 pil的格式转换"""
        return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA))

    @classmethod
    def PIL2cv(cls, image):
        return cv2.cvtColor(np.asarray(image), cv2.COLOR_RGBA2BGR)

    def get_background(self) -> np.ndarray:
        """生成颜色大小的背景图片"""
        background = np.ones((self.size[0], self.size[1], 3), dtype=np.uint8)
        background[:, :, 0] = self.background_color[0]  # 蓝色通道B
        background[:, :, 1] = self.background_color[1]  # 绿色通道G
        background[:, :, 2] = self.background_color[2]  # 红色通道R
        return background

    def get_bounding_rect(self, image: np.ndarray) -> np.ndarray:
        """找到最大轮廓的最小矩形"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 1, 225)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # 找到轮廓中最大面积的然后返回
        max_area = 0
        index = None
        for i, cnt in enumerate(contours):
            x, y, w, h = cv2.boundingRect(cnt)
            area = w * h
            if area > max_area:
                max_area = area
                index = i

        if index is not None:
            x, y, w, h = cv2.boundingRect(contours[index])
            return image[y:y + h, x:x + w, :]
        return image

    @classmethod
    def replace_background_color(cls, image: np.ndarray, color) -> Image.Image:
        """替换背景图片， BGRA格式"""
        image_bg = Image.new('RGB', (image.shape[1], image.shape[0]), color)
        image_with_alpha = cls.cv2PIL(image)
        alpha = image_with_alpha.split()[-1]
        mask = Image.eval(alpha, lambda x: 255 if x > 0 else 0)
        image_bg.paste(image_with_alpha, mask=mask)
        return image_bg

    def transform(self, image: np.ndarray) -> Image.Image:
        background = self.get_background()
        image_bg = self.replace_background_color(image, self.background_color)

        # 找到图像的最小外接矩形
        image = self.get_bounding_rect(self.PIL2cv(image_bg))

        # 放缩比例
        if image.shape[0] / background.shape[0] > image.shape[1] / background.shape[1]:
            image_scale = background.shape[1] * (image.shape[1] / image.shape[0])
        else:
            image_scale = background.shape[1]

        image_resize = imutils.resize(image, int(image_scale * self.scale))
        image_w = image_resize.shape[1]
        image_h = image_resize.shape[0]
        margin_w = int((background.shape[1] - image_w) / 2)
        margin_h = int((background.shape[0] - image_h) / 2)

        image_pil = self.cv2PIL(image_resize)
        background_pil = self.cv2PIL(background)
        background_pil.paste(image_pil, box=(margin_w, margin_h, margin_w + image_w, margin_h + image_h))
        return background_pil

    @classmethod
    def save_filename(cls, image, filename) -> str:
        """保存图片为路径"""
        if isinstance(image, Image.Image):
            image.save(filename)
        elif isinstance(image, np.ndarray):
            cv2.imencode(".png", image)[1].tofile(filename)


def get_image_matting(origin_path):
    """抠图的到图像的mask"""
    result = shop_seg(origin_path)
    mask = result[OutputKeys.MASKS]
    image = cv2.imdecode(np.fromfile(origin_path, dtype=np.uint8, ), 1)
    image_matting = Mask2Image(mask=mask, image=image).transform_png()
    return image_matting


def matting_image_by_path(origin_path, save_path, save_white_path):
    image_matting = get_image_matting(origin_path)
    image_matting_white = ReplaceBackground.replace_background_color(image_matting, color=(255, 255, 255))
    ReplaceBackground.save_filename(image_matting, save_path)
    putCenter(save_path)
    ReplaceBackground.save_filename(image_matting_white, save_white_path)


def image_matting_func():
    image_matting_done_list = [item.stem for item in image_matting_done.iterdir()]
    for file in image_matting_origin.iterdir():
        file_name = file.stem
        save_file_path = pathlib.Path.joinpath(image_matting_done, "%s.png" % file_name)
        save_white_path = pathlib.Path.joinpath(image_matting_done, "%s_white.png" % file_name)
        if save_white_path.stem in image_matting_done_list:
            continue
        else:
            print("matting_action： %s" % file_name)
            try:
                import time
                st = time.time()
                matting_image_by_path(file.as_posix(), save_file_path.as_posix(), save_white_path.as_posix())
                print("matting_action： %s success, duration: %s" % (file_name, str(time.time() - st)))
            except Exception as e:
                print("loop_matting error： %s" % str(e))
                import traceback
                print(traceback.format_exc())


def loop_matting():
    while True:
        try:
            image_matting_func()
        except Exception as e:
            continue


if __name__ == '__main__':
    loop_matting()
